import datetime
import os
import numpy as np
import seaborn as sns
import joblib
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.metrics import f1_score, confusion_matrix, roc_auc_score
from utils.common import data_preprocessing
from utils.log import Logger
import warnings
warnings.filterwarnings('ignore')

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15
pd.set_option('display.max_columns', None)


class AttritionModel:
    def __init__(self, base_path):
        # 1.日志对象
        predict_log_name = 'predict_' + datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        self.logger = Logger(base_path, predict_log_name).get_logger()
        # 2.加载预处理的数据(测试集数据),test数据集在训练时不能使用,专门用来测试模型的效果的,
        # 如果使用test训练的时候进行评估,会造成乐观估计(相当于泄露的未来信息)
        self.data_source = data_preprocessing(os.path.join(base_path, 'data/test.csv'))
        # 3.加载模型
        self.model = joblib.load(os.path.join(base_path, 'model/xgboost_model.pth'))

    def pred_feature_extract(self, feature):
        # 构建特征列表list
        # train_feature_name = ['TotalWorkingYears', 'OverTime', 'MonthlyIncome', 'Age', 'YearsAtCompany',
        #                       'JobLevel', 'YearsInCurrentRole', 'YearsWithCurrManager', 'StockOptionLevel',
        #                       'JobSatisfaction', 'JobInvolvement', 'EnvironmentSatisfaction',
        #                       'YearsSinceLastPromotion',
        #                       'MaritalStatus_Divorced', 'MaritalStatus_Married', 'MaritalStatus_Single',
        #                       'JobRole_Healthcare Representative', 'JobRole_Human Resources',
        #                       'JobRole_Laboratory Technician', 'JobRole_Manager', 'JobRole_Manufacturing Director',
        #                       'JobRole_Research Director', 'JobRole_Research Scientist', 'JobRole_Sales Executive',
        #                       'JobRole_Sales Representative']
        train_feature_name = ['StockOptionLevel', 'JobLevel', 'EnvironmentSatisfaction', 'YearsWithCurrManager', 'OverTime', 'Age', 'WorkLifeBalance', 'JobInvolvement', 'NumCompaniesWorked', 'MonthlyIncome', 'DistanceFromHome', 'Department_Human Resources', 'Department_Research & Development', 'Department_Sales', 'MaritalStatus_Divorced', 'MaritalStatus_Married', 'MaritalStatus_Single']

        train_feature_list = []

        for feature_name in train_feature_name:
            try:
                if feature_name in feature.index:
                    # 直接数值特征
                    value = feature[feature_name]
                    # 特别处理OverTime特征（必须与训练时完全一致）
                    if feature_name == 'OverTime':
                        value = 1 if str(value).lower() == 'yes' else 0
                    train_feature_list.append(float(value))
                elif '_' in feature_name and len(feature_name.split('_')) >= 2:
                    # one-hot 编码特征处理（必须与训练时完全一致）
                    parts = feature_name.split('_', 1)
                    main_feature = parts[0]
                    category_value = parts[1]
                    if main_feature in feature.index:
                        actual_value = str(feature[main_feature])
                        train_feature_list.append(1 if actual_value == category_value else 0)
                    else:
                        train_feature_list.append(0)
                else:
                    train_feature_list.append(0)
            except Exception as e:
                print(f"处理特征 {feature_name} 时出错: {e}")
                train_feature_list.append(0)

        # 创建 DataFrame 并确保数据类型正确
        result_df = pd.DataFrame([train_feature_list], columns=train_feature_name)

        # 确保所有列都是数值类型
        for col in result_df.columns:
            result_df[col] = pd.to_numeric(result_df[col], errors='coerce').fillna(0)

        return result_df

    def evaluate_plot(self, evaluate_df):
        """
        绘制预测结果的图像
        :return:
        """
        fig, axes = plt.subplots(2, 2, figsize=(20, 15))

        # 1. 混淆矩阵
        cm = confusion_matrix(evaluate_df['真实结果'], evaluate_df['预测结果'])
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', ax=axes[0, 0],
                    xticklabels=['未离职', '离职'],
                    yticklabels=['未离职', '离职'])
        axes[0, 0].set_title('混淆矩阵')
        axes[0, 0].set_xlabel('预测结果')
        axes[0, 0].set_ylabel('真实结果')

        # 2. 类别分布对比
        true_counts = evaluate_df['真实结果'].value_counts().sort_index()
        pred_counts = evaluate_df['预测结果'].value_counts().sort_index()

        all_labels = sorted(set(true_counts.index) | set(pred_counts.index))
        true_counts = true_counts.reindex(all_labels, fill_value=0)
        pred_counts = pred_counts.reindex(all_labels, fill_value=0)

        x = np.arange(len(all_labels))
        width = 0.35
        axes[0, 1].bar(x - width / 2, true_counts.values, width, label='真实结果', alpha=0.8)
        axes[0, 1].bar(x + width / 2, pred_counts.values, width, label='预测结果', alpha=0.8)
        axes[0, 1].set_xlabel('类别')
        axes[0, 1].set_ylabel('数量')
        axes[0, 1].set_title('类别分布对比')
        axes[0, 1].set_xticks(x)
        axes[0, 1].set_xticklabels([f'类别{i}' for i in all_labels])
        axes[0, 1].legend()

        # 3. 散点图
        np.random.seed(42)
        noise_true = np.random.normal(0, 0.05, len(evaluate_df))
        noise_pred = np.random.normal(0, 0.05, len(evaluate_df))

        axes[1, 0].scatter(range(len(evaluate_df)), evaluate_df['真实结果'] + noise_true,
                           c='red', alpha=0.6, label='真实结果', s=20)
        axes[1, 0].scatter(range(len(evaluate_df)), evaluate_df['预测结果'] + noise_pred,
                           c='blue', alpha=0.6, label='预测结果', s=20)
        axes[1, 0].set_xlabel('样本索引')
        axes[1, 0].set_ylabel('结果')
        axes[1, 0].set_title('预测结果散点图')
        axes[1, 0].legend()

        # 4. 准确率统计
        correct = (evaluate_df['真实结果'] == evaluate_df['预测结果']).sum()
        total = len(evaluate_df)
        accuracy = correct / total

        axes[1, 1].bar(['准确率'], [accuracy], color='green', alpha=0.7)
        axes[1, 1].set_ylabel('准确率')
        axes[1, 1].set_title(f'模型准确率: {accuracy:.4f}')
        axes[1, 1].set_ylim(0, 1)

        # 添加数值标签
        for i, v in enumerate([accuracy]):
            axes[1, 1].text(i, v + 0.01, f'{v:.4f}', ha='center', va='bottom')

        plt.tight_layout()
        plt.show()

    def model_predict(self):
        """
        模型预测
        :return:
        """
        self.logger.info('===============开始进行预测===============')
        # 初始化评估的list
        evaluate_list = []
        data = self.data_source.copy()
        pred_feature = pd.DataFrame()
        for i in range(len(data)):
            # 解析特征
            pred_feature = self.pred_feature_extract(data.iloc[i, :])
            pred = self.model.predict(pred_feature)[0]
            evaluate_list.append([data.iloc[i, :]['EmployeeNumber'], data.iloc[i, :]['Attrition'], pred])

        # print(pred_feature)
        evaluate_df = pd.DataFrame(evaluate_list, columns=['员工号', '真实结果', '预测结果'])

        # 评估模型
        f1 = f1_score(evaluate_df['真实结果'], evaluate_df['预测结果'])
        ra = roc_auc_score(evaluate_df['真实结果'], evaluate_df['预测结果'])
        print(f"模型F1值: {f1}")
        print(f"模型ROC-AUC值: {ra}")

        # 绘制图像
        self.evaluate_plot(evaluate_df)
        self.logger.info('===============预测结束===============')
        return evaluate_df


if __name__ == '__main__':
    model = AttritionModel('../')
    # feature = model.data_source.iloc[0, :]
    # pred_feature = model.pred_feature_extract(feature)
    # print(pred_feature)
    model.model_predict()
